package city.spring.utils;

import com.baomidou.mybatisplus.core.toolkit.support.SFunction;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Sets;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.springframework.lang.NonNull;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * <h1>树形结构数据处理工具</h1>
 * <p>Entity 结构（使用 String 做主键）：public Entity(String id, String pid, List&lt;Entity&gt; children)</p>
 * <p>Entity 结构（使用 Integer 做主键）：public Entity(Integer id, Integer pid, List&lt;Entity&gt; children)</p>
 *
 * @author HouKunLin
 * @date 2020/5/3 0003 20:18
 */
public class TreeDataUtils {

    /**
     * 自动整理、归类、收集树形结构数据，使 multimap 的 key 只保留所有数据中的顶级数据。
     *
     * @param list             原始树形结构单节点数据
     * @param defaultParentKey 默认的顶级父级ID
     * @param getId            从 Entity 获取当前ID的方法
     * @param getParentId      从 Entity 获取父级ID的方法
     * @param getChildren      从 Entity 获取子级列表的方法
     * @param setChildren      给 Entity 设置子级列表的方法
     * @param <E>              Entity 实体类对象
     * @param <K>              当前对象ID 与 父级对象ID 关联关系的类型，比如：String 、 Integer 类型数据来做关联
     * @return <p>
     * 返回处理后的树形结构信息，返回数据的结构为：<br/>
     * 键-值：list中所有顶级父级ID -> 该顶级父级ID的子级列表，一般情况下顶级父级ID为defaultParentKey
     * </p>
     */
    public static <E, K> ListMultimap<K, E> handlerTreeDatasource(@NonNull List<? extends E> list,
                                                                  K defaultParentKey,
                                                                  Function<E, @Nullable K> getId,
                                                                  Function<E, @Nullable K> getParentId,
                                                                  Function<E, @Nullable List<E>> getChildren,
                                                                  BiConsumer<E, List<E>> setChildren) {
        ListMultimap<K, E> multimap = ArrayListMultimap.create();
        // 获取父级ID，因为顶级id可能为null，因此需要提供一个默认的父级ID处理方式
        Function<K, @Nullable K> getParentKey = (key) -> {
            if (key == null) {
                return defaultParentKey;
            }
            return key;
        };
        // 初步合并树形结构，把所有数据中，同一个父级的数据归类到一个List中存储，存储方式（键-值）：父级ID -> 子级列表
        list.forEach(entity -> multimap.put(getParentKey.apply(getParentId.apply(entity)), entity));
        // 这里必须把所有的key取出来，因为在 handlerTreeChildren 中涉及到 multimap 对象的 key 移除
        HashSet<K> strings = new HashSet<>(multimap.keySet());
        strings.forEach(key -> {
            // 获得父节点 key 的子节点列表
            if (multimap.containsKey(key)) {
                multimap.get(key)
                        .forEach(entity -> handlerTreeChildren(multimap, entity, getId, getParentId, getChildren, setChildren));
            }
        });
        return multimap;
    }

    /**
     * <p>处理树形结构数据的子节点和孙子节点信息。</p>
     * <p>从 multimap 中移除当前节点的子级数据，</p>
     * <p>把 multimap 中当前节点的子级信息存到当前节点对象中，</p>
     * <p>继续处理当前节点的子节点的子节点数据，也就是继续处理当前节点的孙子节点数据。</p>
     *
     * @param multimap    完整的树形结构数据信息
     * @param parent      需要处理的父节点信息
     * @param getId       从 Entity 获取当前ID的方法
     * @param getParentId 从 Entity 获取父级ID的方法
     * @param getChildren 从 Entity 获取子级列表的方法
     * @param setChildren 给 Entity 设置子级列表的方法
     * @param <E>         Entity 实体类对象
     * @param <K>         当前对象ID 与 父级对象ID 关联关系的类型，比如：String 、 Integer 类型数据来做关联
     */
    private static <E, K> void handlerTreeChildren(@NonNull ListMultimap<K, E> multimap,
                                                   E parent,
                                                   Function<E, @Nullable K> getId,
                                                   Function<E, @Nullable K> getParentId,
                                                   Function<E, @Nullable List<E>> getChildren,
                                                   BiConsumer<E, List<E>> setChildren) {
        if (getChildren.apply(parent) != null) {
            // 当前对象的子级列表不为null，表示已经处理过，不需要再次处理
            return;
        }
        K id = getId.apply(parent);
        K pid = getParentId.apply(parent);
        if (id != null && id.equals(pid)) {
            // 不允许当前节点的父级是自己，当前节点的父级是自己将会造成死循环
            return;
        }
        // 从 multimap 中移除当前节点的子级数据
        List<E> entities = multimap.removeAll(id);
        // 把 multimap 中当前节点的子级信息存到当前节点对象中
        setChildren.accept(parent, entities);
        // 继续处理当前节点子级节点的子级数据，也就是处理当前节点的孙子信息
        entities.forEach(entity -> handlerTreeChildren(multimap, entity, getId, getParentId, getChildren, setChildren));
    }

    /**
     * 处理树形结构的父级信息、从数据库中查询出来的数据可能不包含父类信息，因此需要再次查询一次
     *
     * @param <E>              Entity 实体类对象
     * @param <K>              当前对象ID 与 父级对象ID 关联关系的类型，比如：String 、 Integer 类型数据来做关联
     * @param multimap         树形结构数据信息
     * @param defaultParentKey 默认父级键
     * @param lambdaQuery      查询对象
     * @param getId            从 Entity 获取当前ID的方法
     * @param setChildren      给 Entity 设置子级列表的方法
     * @return 树形结构结果 list
     */
    public static <E, K> List<E> loadTreeParents(ListMultimap<K, E> multimap,
                                                 K defaultParentKey,
                                                 LambdaQueryChainWrapper<E> lambdaQuery,
                                                 SFunction<E, @Nullable K> getId,
                                                 BiConsumer<E, List<E>> setChildren) {
        List<E> result = new ArrayList<>();
        if (multimap.isEmpty()) {
            return result;
        }
        Set<K> parentKeys = multimap.keySet();
        if (parentKeys.contains(defaultParentKey)) {
            // 添加默认的父级键信息列表到返回结果中
            result.addAll(multimap.get(defaultParentKey));
            // 这里移除 key 会影响 multimap 的值（同时移除相应的key）
            parentKeys.remove(defaultParentKey);
        }
        if (!parentKeys.isEmpty()) {
            // 从数据库中加载父类信息，到这里 parentKeys 一定至少存在一个 key 信息
            List<E> queryResults = lambdaQuery.in(getId, parentKeys).list();
            queryResults.forEach(entity -> setChildren.accept(entity, multimap.get(getId.apply(entity))));
            result.addAll(queryResults);

            // 数据库中查询到的key信息
            Set<K> queryResultKeys = queryResults.stream()
                    .map(getId)
                    .collect(Collectors.toSet());

            // 返回 在set1中但不在set2中 的元素，也就是返回数据库中没有查询到的key列表
            Sets.SetView<K> notInQueryResultKeys = Sets.difference(parentKeys, queryResultKeys);
            // 数据库不存在这些父级信息，表示这是一个悬空的数据（无法找到父级）
            notInQueryResultKeys.forEach(key -> result.addAll(multimap.get(key)));
        }
        return result;
    }
}
